import typing
import asyncio
from functools import partial

from PyQt6 import QtCore, QtNetwork
from PyQt6.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkReply

RequestFutureAttribute = QtNetwork.QNetworkRequest.Attribute.User


class QAsyncNetworkError(Exception):
    def __init__(self, msg='unknown'):
        super(QAsyncNetworkError, self).__init__()
        self.msg = msg


class QAsyncNetworkManager(QNetworkAccessManager):
    def __init__(self, parent: QtCore.QObject = None, loop: asyncio.AbstractEventLoop = None):
        super(QAsyncNetworkManager, self).__init__(parent)
        self.charset = 'UTF-8'
        self.loop = loop if loop is not None else asyncio.get_event_loop()
        self.finished.connect(self._handle_response)

    def createRequest(self, op: 'QNetworkAccessManager.Operation',
                      request: 'QNetworkRequest',
                      device: typing.Optional[QtCore.QIODevice] = ...) -> 'QNetworkReply':
        reply = super().createRequest(op, request, device)
        future = request.attribute(RequestFutureAttribute, None)
        reply.setAttribute(RequestFutureAttribute, future)
        return reply

    def _handle_response(self, reply: QtNetwork.QNetworkReply):
        er = reply.error()
        future: asyncio.Future = reply.attribute(RequestFutureAttribute)
        try:
            if er == QtNetwork.QNetworkReply.NetworkError.NoError:
                rsp_bytes = reply.readAll().data()
                if future is not None:
                    self.loop.call_soon_threadsafe(future.set_result, rsp_bytes)
                    # future.result(rsp_bytes)
            else:
                error_msg = reply.errorString()
                raise QAsyncNetworkError(error_msg)
        except Exception as e:
            # future.set_exception(e)
            self.loop.call_soon_threadsafe(future.set_exception, e)

    async def http_request(self, url: str, method: str, headers: dict = {}, body_data: QtCore.QByteArray = None,
                           timeout: int = 5):
        q_url = QtCore.QUrl(url)
        q_request = QtNetwork.QNetworkRequest()
        q_request.setUrl(q_url)
        q_request.setTransferTimeout(timeout * 1000)

        for k, v in headers.items():
            q_request.setRawHeader(k.encode(self.charset), v.encode(self.charset))

        data = body_data if body_data is not None and body_data.size() > 0 else None

        f = self.loop.create_future()
        q_request.setAttribute(RequestFutureAttribute, f)

        reply = self.sendCustomRequest(q_request, method.upper().encode(self.charset), data)

        await f

        return f.result()

    async def http_get(self, url: str, timeout: int = 5):
        rsp = await self.http_request(url, 'GET', timeout=timeout)
        return rsp


